# build_knowledge_base.py
import os
from pathlib import Path
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from modelscope import snapshot_download

# === 1. 加载 Java / Kotlin / XML 文件（递归，含根目录）===
print("🔍 正在加载 Android 代码（.java, .kt, .xml）...")
docs = []
code_dir = Path("./android_projects")

if not code_dir.exists():
    raise FileNotFoundError(f"代码目录不存在: {code_dir.resolve()}")

# 支持的扩展名
extensions = ["*.java", "*.kt", "*.xml"]

for ext in extensions:
    for file_path in code_dir.rglob(ext):
        try:
            # 尝试 UTF-8 编码，失败则跳过（避免二进制或乱码文件）
            loader = TextLoader(file_path, encoding="utf-8")
            docs.extend(loader.load())
            print(f"✅ 加载: {file_path.relative_to(code_dir)}")
        except UnicodeDecodeError:
            print(f"⚠️  编码错误，跳过: {file_path.relative_to(code_dir)}")
        except Exception as e:
            print(f"❌ 加载失败，跳过: {file_path.relative_to(code_dir)} - {e}")

print(f"✅ 总共加载 {len(docs)} 个代码/配置文件")

if len(docs) == 0:
    print("❌ 错误：未找到任何 .java / .kt / .xml 文件！请检查 android_projects 目录。")
    exit(1)

# === 2. 分块 ===
print("✂️ 正在分块...")
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=300,
    chunk_overlap=50,
    separators=["\n\n", "\n", "{", "}", ";", " ", ""]
)
splits = text_splitter.split_documents(docs)
print(f"✅ 生成 {len(splits)} 个文本块")

# === 3. 使用 ModelScope 加载 BGE 中文 embedding 模型（国内高速）===
print("🧠 正在从 ModelScope 加载 embedding 模型（bge-small-zh-v1.5）...")
model_dir = snapshot_download(
    'AI-ModelScope/bge-small-zh-v1.5',  # 注意：点号转为 ___
    cache_dir='./models'
)

embedding = HuggingFaceEmbeddings(
    model_name=model_dir,
    model_kwargs={'device': 'cpu'},
    encode_kwargs={'normalize_embeddings': True}
)

# === 4. 构建并保存 Chroma 向量库 ===
print("💾 正在构建向量库...")
vectorstore = Chroma.from_documents(
    documents=splits,
    embedding=embedding,
    persist_directory="./chroma_db"
)
print("✅ 知识库构建完成！数据已保存到 ./chroma_db")